import jieba

from tqdm import tqdm
import os
import torch
from loader import loader4_baidu_senti
import run4_baidu_senti
from nezha.modeling.nezha_senti import NeZhaForSequenceClassification, get_word_indexs, get_word_indexs2
from fish_tool import logs, sys_tool
from fish_tool.ai import torch_tool
from loader.vocab import Vocab
from paddlenlp import Taskflow

jieba.load_userdict(os.path.join(os.path.dirname(__file__), "jiebaDict.txt"))


def inter1():
    config = run4_baidu_senti.TrainConfig()
    vocab = Vocab(config.pre_model_dir)
    data = loader4_baidu_senti.get_data(config, 'test', batch_size=1)
    for batch in data:
        for b in range(batch['input_ids'].size(0)):
            txt = ''.join([vocab.trans_id_to_token(t) for t in batch['input_ids'][b].tolist()])
            logs.print(txt)
            words = list(jieba.cut(txt))
            logs.print(words)
            exit()


def inter2():
    config = run4_baidu_senti.TrainConfig()
    seg = Taskflow('word_segmentation')
    vocab = Vocab(config.pre_model_dir)
    data = loader4_baidu_senti.get_data(config, 'train', batch_size=1)
    i = 0
    for batch in data:
        for b in range(batch['input_ids'].size(0)):
            txt = ''.join([vocab.trans_id_to_token(t) for t in batch['input_ids'][b].tolist()])
            jieba_words = list(jieba.cut(txt))
            baidu_words = seg(txt)
            if jieba_words != baidu_words:
                # logs.print(txt)
                logs.print(f'jieba_words={jieba_words}')
                # logs.print(f'baidu_words={baidu_words}')
            i += 1
        if i > 200:
            break


def show_cut_words(test_num=30):
    config = run4_baidu_senti.TrainConfig()
    vocab = Vocab(config.pre_model_dir)
    data = loader4_baidu_senti.get_data(config, 'test', batch_size=1)
    for i, batch in enumerate(data):
        input_ids = batch['input_ids'].tolist()
        for ids in input_ids:
            tokens = [vocab.trans_id_to_token(t) for t in ids]
            word_indexs = get_word_indexs(tokens)
            tokens_list = [tokens[w[0]: w[-1] + 1] for w in word_indexs]
            words = [''.join(t) for t in tokens_list]
            print(words)
        if i > test_num:
            break


print(get_word_indexs2(['CLS'] + list('这家店就是垃圾。这家店非常垃圾。')))
exit()

seg1 = Taskflow('ner')


seg2 = Taskflow('lexical_analysis')
print(seg1('没有比这更差的酒店了。房间灯光暗淡，空调无法调节，前台服务僵化。'))
print(seg1('这家店就是垃圾。这家店非常垃圾。这家店太垃圾了。'))
print(seg2('没有比这更差的酒店了。房间灯光暗淡，空调无法调节，前台服务僵化。'))
exit()


def xxx():
    type__txt = {}
    config = run4_baidu_senti.TrainConfig()
    vocab = Vocab(config.pre_model_dir)
    data = loader4_baidu_senti.get_data(config, 'test', batch_size=1)
    for batch in data:
        input_ids = batch['input_ids'].tolist()
        for ids in input_ids:
            tokens = [vocab.trans_id_to_token(t) for t in ids]
            txt = ''.join(tokens)
            for word, wtype in seg1(txt):
                type__txt[wtype] = word
    logs.print(type__txt)
    """{'w': '，',
     '世界地区类': '祈年殿',
     '个性特征': '霸道',
     '事件类': '杯具',
     '人物类_实体': '巴适',
     '人物类_概念': '霸道总裁',
     '介词': '跟',
     '介词_方位介词': '之前',
     '代词': '这家',
     '位置方位': '周围',
     '作品类_实体': 'CLS老板人好好，东西好好',
     '作品类_概念': '盗版书',
     '信息资料': '关键',
     '其他角色类': '渔家',
     '副词': '很',
     '助词': '的',
     '医学术语类': '疗程',
     '叹词': '哦',
     '品牌名': '老板',
     '场所类': '店铺',
     '场所类_交通场所': '马路',
     '场景事件': '排',
     '宇宙类': '黑洞',
     '感官特征': '味道',
     '拟声词': '咩',
     '数量词': '几个',
     '时间类': '晚上',
     '时间类_特殊日': '节假日',
     '术语类': '财务',
     '术语类_符号指标类': '价格',
     '物体类': '电子',
     '物体类_化学物质': '甲醛污',
     '生物类_动物': '乌鸡',
     '生物类_植物': '参天大树',
     '疑问词': '啥',
     '组织机构类': 'CLS',
     '组织机构类_企事业单位': '公司',
     '组织机构类_体育组织机构': '火箭',
     '组织机构类_医疗卫生机构': '医院',
     '组织机构类_教育组织机构': '电大',
     '肯定词': '要',
     '药物类_中药': '龟苓膏',
     '词汇用语': 'CLS',
     '连词': '和',
     '饮食类': '水果',
     '饮食类_菜品': '菜',
     '饮食类_饮品': '好酒'
     
     '文化类': '猪',
     '药物类': '巴适',
     '术语类_生物体': '口水',
     '疾病损伤类': '脱皮',
         '修饰词': '小',
      '否定词': '还没有',
    }"""


if __name__ == '__main__':
    xxx()
